#include <cmath>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

class OneEuroFilter {
private:
    double x;
    double dx;
    double te;
    double mincutoff;
    double beta;
    double dcutoff;
    double alpha;
    double dalpha;
    bool initialized;

    double calculate_alpha(double cutoff) {
        double tau = 1.0 / (2 * M_PI * cutoff);
        return 1.0 / (1.0 + tau / te);
    }

public:
    OneEuroFilter(double te, double mincutoff = 1.0, double beta = 0.007, double dcutoff = 0.5)
        : te(te), mincutoff(mincutoff), beta(beta), dcutoff(dcutoff), initialized(false) {
        alpha = calculate_alpha(mincutoff);
        dalpha = calculate_alpha(dcutoff);
    }

    py::array_t<double> filter(py::array_t<double> input, double te) {
        auto input_buf = input.request();
        auto result = py::array_t<double>(input_buf.size);
        auto result_buf = result.request();
        
        double* input_ptr = (double*)input_buf.ptr;
        double* result_ptr = (double*)result_buf.ptr;
        
        size_t num_points = input_buf.size / input_buf.shape[1];
        size_t dim = input_buf.shape[1];
        
        for (size_t i = 0; i < num_points; i++) {
            for (size_t j = 0; j < dim; j++) {
                size_t idx = i * dim + j;
                double x = input_ptr[idx];
                
                if (!initialized) {
                    result_ptr[idx] = x;
                    this->x = x;
                    this->dx = 0;
                } else {
                    double edx = (x - this->x) / te;
                    this->dx = this->dx + (this->dalpha * (edx - this->dx));
                    double cutoff = this->mincutoff + this->beta * std::abs(this->dx);
                    this->alpha = calculate_alpha(cutoff);
                    result_ptr[idx] = this->x + this->alpha * (x - this->x);
                }
            }
            this->x = result_ptr[i * dim];
        }
        
        initialized = true;
        result.resize({num_points, dim});
        return result;
    }
};

PYBIND11_MODULE(one_euro_filter, m) {
    py::class_<OneEuroFilter>(m, "OneEuroFilter")
        .def(py::init<double, double, double, double>(),
             py::arg("te"),
             py::arg("mincutoff") = 1.0,
             py::arg("beta") = 0.007,
             py::arg("dcutoff") = 0.5)
        .def("filter", &OneEuroFilter::filter);
} 